import os
import torch
import time
import loguru
import argparse

import pandas as pd
from   chatglm import ChatGLM3Evaluator

logs_dir = "eval_logs"
logger   = loguru.logger
choices  = ["A", "B", "C", "D"]


def main(args):
    evaluator = ChatGLM3Evaluator(choices=choices, 
                                k=args.ntrain,
                                model_name=args.model_name, 
                                logger=logger,
                            )

    subject_name = args.subject
    data_path    = args.data_path
    if not os.path.exists(logs_dir):
        os.mkdir(logs_dir)
        
    run_date     = time.strftime('%Y-%m-%d_%H-%M-%S',time.localtime(time.time()))
    save_dir     = os.path.join(logs_dir, f"{args.model_name}_{run_date}")
    os.mkdir(save_dir)
    
    val_file_path=os.path.join(data_path + "val" if data_path[-1] == ("/") else data_path+"/val", f'{subject_name}_val.csv')  
    logger.info(f"测评数据文件: {val_file_path}, 结果保存路径: {save_dir}")
    
    val_df = pd.read_csv(val_file_path)
    if args.few_shot:
        dev_file_path = os.path.join(data_path + "val" if data_path[-1] == ("/") else data_path+"/val", f'{subject_name}_val.csv')
        dev_df        = pd.read_csv(dev_file_path)
        correct_ratio = evaluator.eval(subject_name, val_df, dev_df,
                                       few_shot=args.few_shot, 
                                    #    save_dir=save_dir, 
                                       cot = args.cot
                                    )
    else:
        correct_ratio = evaluator.eval(subject_name, val_df, 
                                       few_shot=args.few_shot, 
                                    #    save_dir=save_dir
                                    )
    logger.warning(f'模型{args.model_name}在学科{subject_name}上的精度为: {correct_ratio}')

# python eval.py --model_name chatglm3-6b                            \
#                --few_shot --ntrain 5                               \
#                --data_path /home/yangxianpku/datasets/ceval-exam   \
#                --subject computer_architecture
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ntrain", "-k",  type=int, default=5)                   # 指定few-shot的数量
    parser.add_argument("--openai_key",    type=str, default="")                  # OpenAI Key
    parser.add_argument("--few_shot",      action="store_true")                   # 使用启动few-shot学习
    parser.add_argument("--model_name",    type=str, required=True)               # 模型名称
    parser.add_argument("--data_path",     type=str, required=True)               # 数据文件路径
    parser.add_argument("--cot",           action="store_true")                   # 是否启动思维链模式
    parser.add_argument("--subject", "-s", type=str, default="operating_system")  # 学科主题
    args = parser.parse_args()
    main(args)